#!/usr/bin/env python3
'''
Author: Haoran Peng
Email: gavinsweden@gmail.com
'''
import time
import sys
from copy import deepcopy
import cv2
import numpy as np
import yaml
import os

from cbs_mapf.planner import Planner

import imageio


class Simulator:
    if_draw_path = False
    wait = False
    random_color = True
    if_draw_rect = False

    def __init__(self):
        # Set up a white canvas,size the same as the whole frame
        frame_size = [v for v in RECT_OBSTACLES.values()][0][1]
        self.canvas = np.ones((frame_size[1], frame_size[0], 3), np.uint8) * 255
        # Draw the rectangluar obstacles on canvas
        if self.if_draw_rect:
            self.draw_rect(np.array([np.array(v) for v in RECT_OBSTACLES.values()]))

        # Transform the vertices to be border-filled rectangles
        static_obstacles = self.vertices_to_obsts(RECT_OBSTACLES)

        # Call cbs-mapf to plan
        self.planner = Planner(GRID_SIZE, ROBOT_RADIUS, static_obstacles)
        before = time.time()
        self.path = self.planner.plan(START, GOAL, debug=False)
        after = time.time()
        print('Time elapsed:', "{:.4f}".format(after - before), 'second(s)')
        # Assign each agent a colour
        self.colours = self.assign_colour(len(self.path))

        # Put the path into dictionaries for easier access
        d = dict()

        for i, path in enumerate(self.path):
            if self.if_draw_path:
                self.draw_path(self.canvas, path, i)  # Draw the path on canvas
            d[i] = path
        self.path = d

    '''
    Transform opposite vertices of rectangular obstacles into obstacles
    '''

    @staticmethod
    def vertices_to_obsts(obsts):
        def drawRect(v0, v1):
            o = []
            base = abs(v0[0] - v1[0])
            side = abs(v0[1] - v1[1])
            for xx in range(0, base, 30):
                o.append((v0[0] + xx, v0[1]))
                o.append((v0[0] + xx, v0[1] + side - 1))
            o.append((v0[0] + base, v0[1]))
            o.append((v0[0] + base, v0[1] + side - 1))
            for yy in range(0, side, 30):
                o.append((v0[0], v0[1] + yy))
                o.append((v0[0] + base - 1, v0[1] + yy))
            o.append((v0[0], v0[1] + side))
            o.append((v0[0] + base - 1, v0[1] + side))
            return o

        static_obstacles = []
        for vs in obsts.values():
            static_obstacles.extend(drawRect(vs[0], vs[1]))
        return static_obstacles

    '''
    Randomly generate colours
    '''

    @staticmethod
    def assign_colour(num):
        def colour(x):
            x = hash(str(x + 42))
            return x & 0xFF, (x >> 8) & 0xFF, (x >> 16) & 0xFF

        colours = dict()
        for i in range(num):
            colours[i] = colour(i)
        return colours

    def draw_rect(self, pts_arr: np.ndarray) -> None:
        for pts in pts_arr:
            cv2.rectangle(self.canvas, tuple(pts[0]), tuple(pts[1]), (0, 0, 255), thickness=3)

    def draw_path(self, frame, xys, i):
        for x, y in xys:
            if self.random_color:
                cv2.circle(frame, (int(x), int(y)), ROBOT_RADIUS // 2, self.colours[i], -1)
            else:
                cv2.circle(frame, (int(x), int(y)), ROBOT_RADIUS // 2, (255, 0, 0), -1)

    '''
    Press any key to start.
    Press 'q' to exit.
    '''
    def start(self, gif_path=None):
        frames = []
        cv2.namedWindow('frame', cv2.WINDOW_NORMAL)
        frame_size = [v for v in RECT_OBSTACLES.values()][0][1]
        cv2.resizeWindow('frame', (frame_size[0] + 1, frame_size[1] + 1))
        # 每一帧绘制一张
        try:
            i = 0
            while True:
                # deepcopy防止出现混淆
                frame = deepcopy(self.canvas)
                # 每一帧的每个路径，迭代绘制路径
                for id_ in self.path:
                    x, y = tuple(self.path[id_][i])
                    if self.random_color:
                        cv2.circle(frame, (x, y), ROBOT_RADIUS, self.colours[id_], -1)
                    else:
                        cv2.circle(frame, (x, y), ROBOT_RADIUS, (255, 0, 0), -1)
                frames.append(frame)
                cv2.imshow('frame', frame)
                if self.wait:
                    cv2.waitKey(10)
                i += 1
        except Exception:
            cv2.waitKey(0)
            cv2.destroyAllWindows()
        if gif_path is not None:
            print(f'Frame num: {len(frames)}')
            step = 7
            print([index * step for index in range(len(frames) // step + 1)])
            sub_frames = [frames[index * step] for index in range(len(frames) // step + 1)]
            if len(frames) % step != 0:
                sub_frames.append(frames[-1])
            # 保存gif，步长1秒
            image_count = 0
            for frame in sub_frames:
                cv2.imwrite(f'../fig/image/{image_count}.jpg', frame)
                image_count += 1
            imageio.mimsave(gif_path, sub_frames, 'GIF', duration=1)


def load_scenario(fd):
    with open(fd, 'r') as f:
        global GRID_SIZE, ROBOT_RADIUS, RECT_OBSTACLES, START, GOAL
        data = yaml.load(f, Loader=yaml.FullLoader)
        GRID_SIZE = data['GRID_SIZE']
        ROBOT_RADIUS = data['ROBOT_RADIUS']
        RECT_OBSTACLES = data['RECT_OBSTACLES']
        START = data['START']
        GOAL = data['GOAL']


'''
这里是运行代码过程中可视化的，因为会影响速度所以只显示一秒
'''


def show_pos(pos):
    cv2.namedWindow('frame', cv2.WINDOW_NORMAL)
    cv2.resizeWindow('frame', (1280, 720))
    # 用于方便查看
    frame = np.ones((1080, 1920, 3), np.uint8) * 255
    for x, y in pos:
        cv2.circle(frame, (x, y), ROBOT_RADIUS, (0, 0, 0), -1)
    cv2.imshow('frame', frame)
    cv2.waitKey(1000)
    cv2.destroyAllWindows()


if __name__ == '__main__':
    if len(sys.argv) < 2:
        path = 'question1.yaml'
    else:
        path = sys.argv[1]
    load_scenario(path)
    show_pos(GOAL)
    r = Simulator()
    r.start(gif_path='../fig/' + os.path.splitext(os.path.basename(path))[0] + '.gif')
